"""
  filename      : HLFD
  author        : 13105
  date          : 2025/11/17
  Description   : 
"""
import torch
import torch.nn.functional as F
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import os
from torchvision import transforms
import cv2
from HLNet.models.HFDB import *
# 确保中文显示正常
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 定义图像预处理和后处理函数
def preprocess_image(image_path, img_size=(256, 256)):
    """将图像转换为模型输入格式"""
    transform = transforms.Compose([
        transforms.Resize(img_size),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 归一化到[-1, 1]
    ])

    image = Image.open(image_path).convert('RGB')
    original_size = image.size
    image_tensor = transform(image).unsqueeze(0)  # 添加批次维度
    return image_tensor, original_size

def postprocess_tensor(tensor, original_size=None):
    """将张量转换为可显示的图像格式"""
    # 去除批次维度
    tensor = tensor.squeeze(0)
    # 反归一化
    tensor = tensor * 0.5 + 0.5  # 从[-1,1]转换到[0,1]
    # 转换为numpy数组并调整通道顺序
    image_np = tensor.permute(1, 2, 0).detach().cpu().numpy()
    # 限制在有效范围内
    image_np = np.clip(image_np, 0, 1)
    # 调整大小到原始尺寸（如果需要）
    if original_size:
        image_np = cv2.resize(image_np, original_size)
    # 转换为8位整数
    return (image_np * 255).astype(np.uint8)

def visualize_features(input_img, low_feat, high_feat, output_img, save_path=None):
    """可视化输入、低频特征、高频特征和输出"""
    fig, axes = plt.subplots(2, 2, figsize=(12, 10))

    axes[0, 0].imshow(input_img)
    axes[0, 0].set_title('输入图像')
    axes[0, 0].axis('off')

    axes[0, 1].imshow(low_feat)
    axes[0, 1].set_title('低频特征（上采样后）')
    axes[0, 1].axis('off')

    axes[1, 0].imshow(high_feat)
    axes[1, 0].set_title('高频特征')
    axes[1, 0].axis('off')

    axes[1, 1].imshow(output_img)
    axes[1, 1].set_title('模型输出')
    axes[1, 1].axis('off')

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"可视化结果已保存至: {save_path}")

    plt.show()

def test_hpb_module(image_path, n_feats=64, save_dir='results'):
    """测试HPB模块并可视化中间过程"""
    # 创建保存目录
    os.makedirs(save_dir, exist_ok=True)

    # 加载并预处理图像
    img_tensor, original_size = preprocess_image(image_path)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    img_tensor = img_tensor.to(device)

    # 如果输入是3通道，需要转换为n_feats通道
    if img_tensor.size(1) != n_feats:
        conv_in = torch.nn.Conv2d(3, n_feats, kernel_size=3, padding=1).to(device)
        img_feats = conv_in(img_tensor)
    else:
        img_feats = img_tensor

    # 初始化HPB模块
    hpb_module = HPB(n_feats, 'haar').to(device)

    # 为了获取中间特征，我们需要修改HPB的forward方法
    # 这里使用钩子函数来捕获中间特征
    low_feat = None
    high_feat = None

    def hook_fn_low(module, input, output):
        nonlocal low_feat
        # 上采样到原始尺寸
        low_feat = F.interpolate(output, size=img_feats.size()[-2:],
                                mode='bilinear', align_corners=True)

    def hook_fn_high(module, input, output):
        nonlocal high_feat
        high_feat = output

    # 注册钩子
    handle_low = hpb_module.unet.register_forward_hook(hook_fn_low)
    handle_high = hpb_module.dense.register_forward_hook(hook_fn_high)

    # 前向传播
    output = hpb_module(img_feats)

    # 移除钩子
    handle_low.remove()
    handle_high.remove()

    # 如果特征通道数不是3，转换为3通道以便可视化
    def convert_to_3ch(feat_tensor):
        if feat_tensor.size(1) != 3:
            # 简单地取前3个通道，或者通过1x1卷积转换
            conv_out = torch.nn.Conv2d(feat_tensor.size(1), 3, kernel_size=1).to(device)
            return conv_out(feat_tensor)
        return feat_tensor

    # 处理特征以便可视化
    low_feat_3ch = convert_to_3ch(low_feat)
    high_feat_3ch = convert_to_3ch(high_feat)
    output_3ch = convert_to_3ch(output)

    # 后处理所有图像
    input_img = postprocess_tensor(img_tensor, original_size)
    low_img = postprocess_tensor(low_feat_3ch, original_size)
    high_img = postprocess_tensor(high_feat_3ch, original_size)
    output_img = postprocess_tensor(output_3ch, original_size)

    # 可视化
    save_path = os.path.join(save_dir, 'hpb_visualization.png')
    visualize_features(input_img, low_img, high_img, output_img, save_path)

    return input_img, low_img, high_img, output_img

if __name__ == "__main__":
    # 替换为你的测试图像路径
    test_image_path = "F:\BaiduNetdiskDownload\SWXD_Data\processed_slices\defective_regions\A_bam5_slice_3.png"  # 请替换为实际图像路径

    # 检查图像文件是否存在
    if not os.path.exists(test_image_path):
        print(f"错误：找不到图像文件 {test_image_path}")
        print("请替换为实际的图像路径")
    else:
        # 运行测试
        test_hpb_module(
            image_path=test_image_path,
            n_feats=64,  # 特征通道数，与模型一致
            save_dir='hpb_results'  # 结果保存目录
        )